Skip to content

feat(models): add AR value head and GAE track wiring - #1

Closed
yhl48 wants to merge 4 commits into
feat/gae-advantagefrom
feat/ar-value-head
Closed

feat(models): add AR value head and GAE track wiring#1
yhl48 wants to merge 4 commits into
feat/gae-advantagefrom
feat/ar-value-head

Conversation

@yhl48

@yhl48 yhl48 commented Jul 25, 2026

Copy link
Copy Markdown
Owner

Summary

Part 2/3 for Tencent-Hunyuan#86 (PPO value critic + GAE on AR models). Stacks on Tencent-Hunyuan#254 (feat/gae-advantage).

  • Adds ValueHead and ARReplayOutput; optional use_value_head on Qwen3 attaches a per-token critic during replay.
  • Extends TextSegment with packed values, returns, and token_advantages.
  • Adds scatter_terminal_rewards (scalar outcome reward → last response token) and RolloutTrack.compute_gae_advantages (per-sample GAE via cu_seqlens, no cross-trajectory leak).
  • Qwen3 replay returns critic values from the same hidden states as log-probs; fails fast if return_values=True without a value head.

Not wired into trainer or PPO loss yet — that is PR 3.

Related Issue

Part of Tencent-Hunyuan#86 (roadmap Tencent-Hunyuan#25). Stacks on Tencent-Hunyuan#254.

Test Plan

CPU (Python 3.12+ env):

pytest tests/types/test_advantages_gae.py tests/types/test_rollout_track_gae.py tests/models/test_value_head.py -v

Locally verified on Python 3.9 for the track/GAE tests above; full project pytest needs 3.12+.

GPU end-to-end (Qwen3 + use_value_head=True + replay values): not run in CI here; follow-up in PR 3 recipe.

Compatibility / Risk

  • New optional config flag use_value_head: bool = False on Qwen3 — default off, no checkpoint/API break for existing GRPO runs.
  • Value head weights are not in base HF checkpoints; training from scratch when enabled.

Reviewer Notes

Checklist

  • I reviewed the changed code and removed unrelated/generated artifacts.
  • I updated tests, docs, and configs where needed, or explained why not.

yhl48 added 4 commits July 25, 2026 17:55
Wire per-token critic values through Qwen3 replay and compute GAE
advantages on RolloutTrack for the PPO critic path (issue Tencent-Hunyuan#86, part 2/3).
Keep transformers imports inside from_config / _packed_replay_supported
as in upstream; only UniRL imports (ValueHead, ARReplayOutput) stay top-level.
Call _require_value_head_for_replay at the start of _replay_aware_forward
so return_values=True fails fast before the transformer forward, not only
from Qwen3ARStage.replay().
Extend ReplayResult with optional per-token values for PPO/GAE and
replace ARReplayOutput in Qwen3 replay with the shared type.
@yhl48

yhl48 commented Jul 26, 2026

Copy link
Copy Markdown
Owner Author

Closing in favor of upstream Tencent-Hunyuan#256 (rebased onto main after Tencent-Hunyuan#254 merged).

@yhl48 yhl48 closed this Jul 26, 2026
leviking98z-rgb added a commit that referenced this pull request Aug 1, 2026
…cent-Hunyuan#214)

* types: add Sample/Part endomorphism types with sample_id lineage (LIN-446)

Squash merge of origin/LIN-446/main into LIN-444/main. Combines:
- add Sample/Part endomorphism types
- derive Sample/Part lineage from the sample_ids path

* rollout: convert the rollout engine to Sample → Sample (LIN-454)

Replace the RolloutReq/RolloutResp request/response pair with the unified
recursive Sample/Part types across the dedicated rollout engines —
generate: Sample → Sample, in place, no bridge. The model pipelines, the
trainside engine, the trainers, reward, and the train stack stay on the old
types for a later pass, so this branch is intentionally not runnable as a
full training loop yet.

Engines converted:
- sglang (AR text/vlm), sglang_diffusion (image), vllm_omni (sd3 / qwen_image
  / hv15 single-DiT + HI3 multi-stage), composed (PE AR→diffusion chain).
- base.py ABC flipped to generate(self, sample: Sample) -> Sample;
  engine/__init__ chunked_engine_generate_req -> chunked_engine_generate
  (split-by-root -> regroup -> concat, the Sample analogue of slice-by-index).

Mechanics:
- Pre-forked gen shells are filled by sampling_params type; positional lineage
  replaces parent_track; σ pinned onto DiffusionSamplingParams.sigmas; the x_T
  noise key is derived from the path-id lineage (OD-2); stage_config moved to
  Part.control.
- Multi-input multimodal (chained input Parts) is gated with clear deferral
  raises — a deferred non-goal of this pass.

P0 type fields: Part.control / conditions / fill,
DiffusionSamplingParams.sigmas / init_noise_latent_shape,
LatentSegment.initial_latents.

Verification (rollout-only, no training):
- scripts/check_sample_roundtrip.py structural oracle (5/5 contracts).
- Real GPU rollout smokes pass for every converted engine and modality class:
  AR (sglang/Qwen3), diffusion (vllm_omni sd3, sglang_diffusion sd3), and the
  multi-stage composed PE chain (AR→diffusion).

* rollout: support multi-input multimodal (IT2I + siblings) (LIN-454)

Wire the previously-gated image+text / cot_text modalities (hi3_it2i,
hi3_i2t, the sglang vlm path, and hi3_dit_recaption) onto the
generate(Sample) -> Sample boundary. The type layer already supported it,
so this is mostly un-gating: a second input rides as a chained input Part
via the new Part.input_child(primitive) (branch-1 child, no sampling_params)
so only the head stays a root, and Sample.conditioning() surfaces both
primitives in turn order. Adapters locate the chained input by primitive
type (image_input_part / new cot_text_from_sample); hi3_dit_recaption gets
the ported cot_text build.

Also fixes sglang text.py build_response to preserve intermediate input
Parts (return [*parts[:-1], filled] instead of [parts[0], filled]) so a
multi-input chain's gen Part keeps a valid parent.

Verified engine-only with hand-built Samples: oracle check_sample_roundtrip
7/7 (incl. multi-input image-chain + cot_text fixtures), a real Qwen2.5-VL
VLM run, and a full HunyuanImage-3 80B IT2I GPU run (image+text -> AR
recaption -> edited image, chain filled).

* LIN-480: migrate trainers/train-stack/logging onto the Sample/Part model

Migrate the consumer side (all five trainers + train stack + wandb logging) from
RolloutReq/RolloutResp/RolloutTrack to the Sample/Part endomorphism model
(LIN-446) — the counterpart to LIN-454's rollout-engine conversion. Trainers now
build request Samples (Part.input + .fork) and consume response Samples (frontier
= parts[-1], stages located by sampling_params type); GRPO advantage /
reward-propagation / token-balancing run through the methods already on Part/Sample.

Scope: all trainers incl. it2i (multi-input via Part.input_child) and the HI3
unified two-engine path. Out of scope / unchanged: the reward service
(score_and_attach migrates on a separate branch — its (*, req, track) signature
does not yet match the trainers' score_and_attach(sample) call, so the train tail
is gated on that branch), the trainside engine, and the rollout-engine internals.

Highlights:
- Sample/Part navigation on the type: root_group_ids, gen_parts/gen_part/
  gen_part_index, with_parts, and tree-aware slice/select (a request Sample
  DP-scatters by prompt-tree, not by the parts list).
- Per-trainer request builders + the HI3 two-engine _run_rollout_one stitching
  (AR recaption -> DiT image); the DiT x_T noise key is re-rooted from the
  globally-unique lineage so it stays unique across dp>1 replicas.
- UnifiedModelTrainStack.train_track scatters the [input, ar, image] lineage as a
  unit so both stages shard at the same prompt boundaries.
- wandb on Sample/parts: diffusion stage named "diffusion" consistently;
  num_samples reports the generated-sample count; zero-std groups bucket by the
  advantage grouping (advantage_group_ids).
- disable_driver_xt flag on DiffusionSamplingParams (restores the unified DiT
  driver-x_T escape hatch).

Verification: scripts/check_sample_roundtrip.py (14 contracts) + SD3 (vllm-omni),
AR (sglang), and HI3 two-engine generation smokes pass on GPU. The dp>1 cluster
is contract-tested + A2 engine-verified; its multinode end-to-end behavior is not
yet verified, and the full train loop awaits the separate reward migration.

* LIN-479: model pipelines + trainside engine onto the Sample/Part model

Convert the model bundle's rollout boundary from RolloutReq/RolloutResp to the
recursive Sample/Part types — the model-bundle half of the refactor, pairing with
LIN-480 (trainers / train-stack):

- SD3Pipeline.generate / Qwen3Pipeline.generate take a request Sample (pre-forked
  gen frontier) and fill it; a factored _conditions_for() is the single encode
  path shared by rollout and trainer-side replay (re-encode; Part.conditions left
  empty — no cache).
- TrainsideRolloutEngine.generate -> Sample -> Sample: _ensure_sample_sigmas pins
  sigma onto the gen part's DiffusionSamplingParams; Part.slice/concat chunking.
- NoiseRecipe.from_sample() — Sample-shaped x_T builder keyed on the gen part's
  path-id lineage.
- check_sample_roundtrip.py +2 contracts (from_sample lineage, generate-fills-
  frontier); trainside_{sd3,ar}_smoke.py real-GPU rollout+replay smokes.

Validated on H20 (pod unirl-gz-2): CPU oracle green; SD3 rollout->replay
mean|dlogp|=0 (ratio~1); Qwen3 AR replay self-consistency 0 (the
old_logp_source=replay contract — the in-process autoregress bf16 record vs
replay fp32 differs ~3-5, harmless since trainside recipes use replay for
old_logp). Stages (diffusion.py/ar.py) untouched.

* reward: migrate score_and_attach to Sample → Sample (LIN-481)

Convert the reward adapter from RolloutReq/RolloutTrack to the Sample/Part
endomorphism model — the consumer-side counterpart LIN-480 left gated. The
trainers already call score_and_attach(sample); this matches that contract.

- RewardService.score_and_attach(*, req, track) -> RolloutTrack becomes
  score_and_attach(self, sample: Sample) -> Sample: score the frontier Part,
  build input primitives from Sample.conditioning() (nearest-ancestor caption),
  generated from frontier.primitive keyed by the backend preferred_input_kind,
  and root-sourced prompt metadata via the new Sample.root_metadata.
- AR truncation/overlong shaping preserved verbatim, re-gated on the scored
  frontier being ARSamplingParams.
- Delete the request/track expansion + metadata-normalization machinery
  (_KIND_TO_KEY, _normalize_prompt_metadata, _build_request_for_track): the
  Sample is row-aligned and DP-scatters by prompt-tree, nothing to reconcile.
- Add primitive_modality_key + Sample.root_metadata; refresh reward
  docstrings/README to Sample/Part.
- Add scripts/check_reward_roundtrip.py structural oracle (6 contracts;
  passes on-GPU-pod with real torch).

* fix(sd3): carry encoded conditions on the trainside gen Part for train-stack replay

SD3Pipeline.generate left Part.conditions empty (a "replay re-encodes"
assumption), but the train stack's prepare_segment reads part.conditions
directly with no re-encode wiring — so trainside SD3 training crashed in
SD3Conditions.from_dict (text=None). Every sibling diffusion pipeline and
the SD3 sglang_diffusion adapter already populate conditions via
fill(conditions=...); bring trainside SD3 in line.

Surfaced by the SD3 trainside e2e (trainside_sd3_smoke bypassed this by
calling stage.replay directly, never going through train_track).

* fix(vllm-omni): tile prompts to the gen fan-out for served diffusion rollout

texts_from_sample asserted prompt count == gen sample count, but the
request keeps the input Part un-fanned (one prompt per group) while the
gen shell fans out samples_per_prompt via Sample.fork — so SD3/qwen-image
served rollout crashed on rollout 1 with "prompt count 6 != gen sample
count 96". vllm_omni runs num_outputs_per_prompt=1 (one output per gen
sample), so tile each prompt across its group-by-parent-contiguous gen
siblings. A 1:1 request (no fan-out) is unchanged.

Surfaced by the SD3 sd3_vllmomni e2e (served LIN-454 path).

* fix(sglang): allow unirl weight-sync classes through sglang 0.5.12 SafeUnpickler

sglang 0.5.12's SafeUnpickler (CVE-2025-10164 guard) runs in the engine
process and blocked the colocate tensor weight-sync: update_weights_from_tensor
deserialization rejected unirl's vendored FlattenedTensorBucket /
_rebuild_cuda_tensor_modified (both under the unirl. namespace) with
"Blocked unsafe class loading". Register "unirl." on sglang's server-side
allowlist in SGLangRolloutEngine.__init__ (the tp_size=1/use_ray=False
scheduler runs in this process), mirroring what unirl's own SafeUnpickler
already permits.

Surfaced by the Qwen3 DRPO AR e2e (qwen3_drpo_4b_base_dapo_sglang): the
full vertical (sglang gen -> MathVerify -> DRPO -> FSDP train) ran and
crashed only at the first weight-sync.

* fix(sglang): use native serializer for LLM weight-sync (sglang 0.5.12 compat)

The prior allowlist patch (a1f33d1) registered "unirl." on the engine
process's SafeUnpickler, but sglang's SRT scheduler deserializes the
weight bucket in a SEPARATE spawned subprocess with a fresh allowlist, so
the patch never took (same "Blocked unsafe class loading" recurred).

Instead, serialize with sglang's NATIVE FlattenedTensorBucket /
MultiprocessingSerializer / monkey_patch_torch_reductions (verified
identical device-UUID IPC mapping) so the wire references sglang's own
allowlisted classes — matching what the sglang_diffusion engine already
does. Reverts the ineffective __init__ allowlist patch.

Surfaced by the Qwen3 DRPO AR e2e: full vertical ran, crashed only at the
initial colocate weight-sync into SRT.

* LIN-495: migrate all model bundles onto the Sample → Sample rollout shape

Convert every model pipeline's generate() from RolloutReq → RolloutResp to the
recursive Sample/Part endomorphism, completing the model-bundle half of the
rollout refactor (pairs with the earlier LIN-454/479/480/481 engine/trainer/
train-stack/reward migrations), then retire the legacy types.

- Base contract: Pipeline.generate(self, sample: Sample) -> Sample; lift the
  request-side Sample readers out of the rollout layer into types/sample_ops.py
  so pipelines import them without a models→rollout inversion.
- 7 Tier-1 diffusion pipelines (qwen_image, z_image, wan21, wan22, hunyuan_video,
  hunyuan_video15, ltx2): read prompt via sample.conditioning(), gen params off
  the frontier Part, fill the frontier shell; factor _conditions_for() shared by
  rollout + trainer-side replay. i2v image arrives via the input_child chain.
- qwen_vl + flux2_klein (AR/edit with one extra input modality).
- pe, bagel, hunyuan_image3 (multi-stage 2-gen-part flows): PEPipeline consumes
  the pre-forked [input, ar_shell, diff_shell] Sample and drives the child
  pipelines (mirrors the migrated ComposedRolloutEngine); bagel dispatcher reads
  task from parts[0].control and fills one or both gen Parts; HI3 dispatcher + 5
  modes/ (incl. two-part t2ti).
- Retire unirl/types/rollout_req.py + rollout_resp.py; relocate PrimitiveValue to
  types/primitives.py; drop the dead NoiseRecipe.from_rollout_req; sweep doc
  breadcrumbs across conditions/algorithms/sde/engine docs.

All pipelines carry conditions=<conds>.to_dict() on frontier.fill — the train
stack reads Part.conditions (GRPO/FlowGRPO re-type via conditions_cls.from_dict);
it does NOT re-encode. (Caught by E2E pe-trainside on unirl-gz-2: an earlier
revision left conditions empty and crashed at the first AR train step with
"Qwen3ARConditions.from_dict: expected d['prompt'] ... got None".)

Validated on TaiJi unirl-gz-2 (8×H20, GZ): compileall green; sd3-trainside reward
grew 0.74→0.79 with ratio=1.0 / advantage_std≈0.49; pe-trainside (SD3+Qwen3-0.6B)
runs the full 3-part chain with both AR + diffusion stacks training.

Adds scripts/trainside_{qwen_image,qwen_vl,pe}_smoke.py (tier-representative GPU
rollout→replay smokes).

* LIN-503: agent-trajectory multi-conditioning — role-aware layer + multi-turn encode (both backends)

Squash of LIN-503/main (8 commits) onto LIN-444/main. Takes the rollout Sample from
single-turn-only to full multi-turn, multi-modal conditioning end-to-end.

- Role-aware trajectory layer on Sample/Part (Phase 1/2): Turn, Part.role /
  resolved_role, turns(), text_conditioning() / vision_conditioning() fail-loud
  renderers, replace_frontier / with_filled_frontier write-back; diffusion/omni
  engines wired; sample_ops retired.
- gap C — sglang AR engine encode consumes the trajectory: text.py / vlm.py build the
  chat conversation from text_/vision_conditioning() via a pure transpose + de-expand
  util (rollout/engine/sglang/utils/conversations.py); resolve_sampling fan-out fixed
  to the last-fork branch (parts[-2]).
- Trainside conjugate — qwen3 / qwen_vl in-process encoders consume turns() via
  models/types/conversations.py (transpose, no de-expand, inline-PIL VLM fusion);
  stale "_conditions_for re-tokenize" docstrings + trainside smokes fixed.
- Behavior-preserving on single-turn (byte-identical); parity-safe by construction
  (replay teacher-forces over the stored conditions["prompt"], not a re-encode).
- Verification: CPU oracles (engine 8/8, trainside 6/6, layer 7/7) + 4 GPU multi-turn
  smokes (sglang/trainside × text/VLM) green on Qwen3-4B-Instruct and
  Qwen2.5-VL-3B-Instruct — the captured prompt carries user→assistant→tool (+ image),
  replay ratio=1.

* feat(rollout): agentic tool-loop rollout — AgentLoop + ToolEnvironment (LIN-492)

Add the agentic (multi-turn tool) rollout layer over the Sample/Part model, on top of
LIN-503 multi-turn conditioning. Scope is ROLLOUT only — trainer wiring, reward, and a
dataset for agentic GRPO are a follow-up (the trainer is still single-turn).

- AgentLoop (unirl/rollout/loop/): environment-driven synchronous loop —
  fork -> generate -> env.step -> observe, bounded by max_turns. Structural
  Environment / RolloutEnginePort protocols; the existing SGLang engine satisfies the
  port unchanged.
- ToolEnvironment + tools/ (Tool, CalculatorTool): parse <tool_call>, run a registered
  tool, feed the result back as a role="tool" observation, stop on a final answer (or
  max_turns). Safe ast-based calculator (no eval).
- Sample.observe(role="tool"): append the world-response as a mask-0 input Part tagged
  for LIN-503 role-aware rendering (<tool_response>). Additive — no existing code path
  changed; the whole merge is 1648 insertions, 0 deletions.

Verified:
- CPU oracles: scripts/tool_env_smoke.py (9 checks), scripts/rollout_loop_smoke.py (6),
  scripts/check_sample_roundtrip.py (16, regression) — all green.
- GPU (Qwen3-4B-Instruct): single-turn tool call (tool_env_ar_smoke); the closed loop —
  model calls the calculator, sees the exact result, and answers (tool_env_ar_loop_smoke);
  and a 72-trajectory reliability run at 97.2% correct with clean termination
  (tool_env_ar_reliability).

Known follow-ups: trainer integration (multi-gen-Part train step + reward + dataset),
per-sample termination for n>1 heterogeneous batches, more tools.

* LIN-499: async per-group rollout-engine interface (squash-merged onto agentic)

Squash-merge of LIN-499/main: the async per-group engine contract (agenerate core + sync generate facade + abort/pause/resume), all five engines migrated, Part.weight_version provenance, DevicePool.worker_max_concurrency knob, and the CPU async-contract test suite (11 tests pass on the merged tree).

Conflict in types/sample.py resolved by keeping BOTH agentic's Part.role (LIN-503) and LIN-499's Part.weight_version — adjacent new fields, not mutually exclusive.

* feat(rollout): agentic multi-turn rollout engine (rank-0 pull coordinator)

AgenticRolloutEngine drives multi-turn (tool-use) rollout: a rank-0 coordinator
(BROADCAST+RANK_ZERO, the NCCLWeightSync pattern) over a DP slab of per-worker
persistent drain loops. Each worker runs one run_until_complete that pulls
single-trajectory tasks from rank 0 and runs them as multi-turn agent loops on its
inner engine's event loop (continuous-batched via the inner backend's semaphore).
generate returns a flat List[Sample] of variable-depth trajectories; the GRPO
n-group is recovered by bucketing on the shared prompt root id.

- base: widen generate to RolloutOutput = Sample | List[Sample] (single-turn
  returns Sample, agentic returns List[Sample]); per-turn seams stay -> Sample.
- engine/agentic: coordinator (set_workers/generate/next_task) + per-worker drain
  (run_drain/_drain/_run_one/_pull); _run_coro delegates to the inner so the drain
  and weight-sync verbs share one lock (the quiesce boundary); a per-worker
  trajectory cap distinct from the backend request semaphore; lifecycle/weight-sync
  verbs delegate to the inner engine.
- loop: add async astep to the Environment protocol + ToolEnvironment (non-blocking
  tool boundary); make step re-entrant (turn derived from the sample, not a mutable
  counter) so one env serves concurrent trajectories; add agenerate to
  RolloutEnginePort.
- tests: CPU contract tests (ragged List[Sample] + bucket-by-root, cap saturation,
  two distinct bounds, pull load-balancing, failure isolation, FIFO next_task) +
  astep parity / re-entrancy / slow-tool-yields.
- scripts: multi-worker GPU smoke. Validated on a 2-worker Qwen3-8B H20 slab:
  ragged List[Sample] with correct tool-use answers, and cross-worker rank-0
  aggregation proven safe (gen segments produced on 2 distinct workers, all spans
  plasma-backed in Ray's global object store, 48 TensorRefs hydrate from the driver)
  — risk #1 closed, no run_drain materialization needed.

LIN-522

* feat(agentic): multi-turn agentic RL stack — AgenticTrainer, deep-research tools, ALFWorld

Land the full multi-turn agentic RL capability on top of AgenticRolloutEngine.

AgenticTrainer (unirl/trainer/agentic.py)
  GRPO over variable-depth trajectory lists: group-relative advantage across the n
  siblings of a prompt, every assistant turn concatenated into ONE on-policy
  train_track step (ratio ~ 1). The per-trajectory reward step (_rewards_and_groups)
  is overridable so tasks swap only the reward SOURCE.

Deep-research task (answer-graded)
  M1 calculator + MathVerify and M2 search/visit tools + LLM-judge reward, with
  recipes and the train_deep_research entrypoint.

ALFWorld baseline (AgenticEnvTrainer, env-sourced reward)
  Stateful per-trajectory ALFWorld adapter, an engine hook that attaches the
  environment's terminal-success return to the last turn, admissible-action snapping
  to avoid TextWorld PDDL crashes, and crash-exclusion (NaN reward -> neutral, zero
  advantage) so env-bug crashes cannot corrupt the GRPO gradient. Verified on-pod:
  task-success reward rises ~0.53 -> ~0.85 (peak 0.94), on-policy, on a fixed 8-game
  set (genuinely multi-turn: up to 10 env steps per trajectory, n=8 GRPO group).

Testing: CPU contract tests cover the plumbing (reward attach, group advantages, env
adapter, train_step assembly). The GPU end-to-end path is validated empirically (the
rising curve), not yet in CI.

* feat(rollout): stateful, decoupled tools for agentic RL (LIN-533)

Add a StatefulTool seam so tools can hold per-trajectory state across
turns, with guaranteed teardown and a first out-of-process tool.

- StatefulTool(Tool): session_start -> execute_session -> session_end,
  keyed by a session id carried in the root control bag. ToolEnvironment
  dispatches on isinstance, so the stateless Tool path is unchanged.
- ToolEnvironment: reset mints a uuid4 session id per stateful tool and
  stamps it into control["tool_sessions"] via _part_with_field (returns
  the request unchanged when there are no stateful tools); step/_run
  dispatch execute_session; async aclose ends sessions in the executor,
  swallowing errors.
- AgenticRolloutEngine._run_one: finally-hook calls env.aclose on every
  path (success, crash, abort), duck-typed via getattr and wrapped so it
  can never re-raise into the drain.
- AlfworldEnv.aclose: reclaim the episode + pooled template — fixes a
  leak when a trajectory dies in the engine between turns.
- SandboxTool: persistent per-session Python REPL subprocess (lazy spawn
  in the executor, select-based timeout, killed on session_end);
  validates cross-turn variable reuse.
- Tests: StatefulTool lifecycle + teardown on success and forced
  exception, SandboxTool REPL, and an ALFWorld leak regression.

* feat(agentic): partial rollout — colocate + async trainers (LIN-531)

Over-generate, interrupt generation at a turn boundary on weight sync, commit the fast complete GRPO groups, carry/drop the slow tail.

- Engine: submit/poll/abort/drained over a background buffered drain; turn-boundary checkpoint-and-resume in _run_one.
- Colocate driver: AgenticPartialTrainer / AgenticEnvPartialTrainer — keeps all GPUs on generation while cutting the straggler tail.
- Disaggregated driver: AsyncAgenticTrainer / AsyncAgenticEnvTrainer.
- Tail policy: carry (resumable tool envs) vs drop (stateful envs, ALFWorld).
- Shared _advantage_train_and_log extracted from the barrier train_step (barrier path preserved); entrypoints + ALFWorld/deep-research recipes; per-rollout turn-histogram + committed/dropped instrumentation.

ALFWorld speed study: colocate-partial beats the barrier only with group-level depth variance — Qwen3-8B ~18% faster, Qwen3-0.6B 27% slower.

Integrates with LIN-533 teardown: the _run_one finally aclose now also fires on the abort/checkpoint path, releasing carried trajectories' episodes/sessions.

* feat(deep-research): AReaL tongyi_deepresearch M2 (search + visit + judge)

Reproduces AReaL's tongyi_deepresearch deep-research agent in UniRL's
agentic stack (sync AgenticTrainer + AgenticRolloutEngine):

- Verbatim Tongyi SYSTEM_PROMPT for the M2 recipe (deep_research_search_judge)
- Hardened SearchTool: serper + serpapi providers, retries/backoff
- Hardened VisitTool toward AReaL tool_visit.py: Jina read retries,
  structured evidence/summary extraction, content truncation
- Robust LLM-judge verdict parsing (negative-first regex; fixes "not correct")
- Cross-config invariant guards (AgenticTrainer + AgenticRolloutEngine)
- prepare_asearcher: default ASearcherBase35k split, streaming load
- Committed GPU training smoke (scripts/train_deep_research_smoke.py)

Validated live: reward climbed 0.22 -> 0.50 over 19 rollouts (Qwen3-1.7B
policy, Qwen2.5-72B judge, batch 128, n=8) before an infra-driven collapse.

Follow-ups (not in this change): multi-turn token-recording reconciliation
(Miles-style trim/accumulate) and a trajectory context-token cap for
long-horizon runs. The shipped recipe defaults to max_turns=8 (safe).

* test(deep-research): adversarial LLM-judge verdict-parse cases

CPU unit test for _parse_verdict guarding the negative-first parse against
the "incorrect" ⊃ "correct" substring trap and "not correct" / "wrong"
phrasings the prior substring test misread. 20 cases; the coverage Phase C
intended but the M2 squash omitted. Verified: 20 passed on agentic@ce13a48d.

* refactor(types): grouping as id projections + is_gen kind signal

Collapse the grouping/lineage duplication on the Sample/Part model:

- sample_id: add ancestor_id(sid, depth) — grouping labels are id-prefix
  projections. root_group_ids and Sample.split project directly; the
  _root_groups_per_part walk (re-deriving what __post_init__ validated)
  is deleted.
- compute_advantages: the group_ids label-list override becomes
  group_layer (None = immediate parent, 0 = root prompt; PE
  diffusion_group_scope="prompt" now passes group_layer=0 instead of
  threading root_group_ids labels). scope stays the normalization mode,
  with the historical global branch (unbiased std) kept verbatim for
  bit-parity with the shipped adv_normalization_scope: global recipes.
- delete Part.advantage_group_ids (its only reader, the zero-std wandb
  metrics, buckets by sibling group_ids) and the dead Part.split.
- kind signal: fork requires sampling_params (a paramless gen shell is
  unrepresentable); Part.is_gen is the single predicate behind
  gen_parts / resolved_role / the wandb gen filter. Ids stay pure
  addresses; role stays presentation-only.
- contracts: check_group_layer_advantages pins per-layer GRPO values,
  the unbiased-std global formula, and fail-loud on paramless fork and
  negative group_layer.

* refactor(rollout): make sync and async paths engine-native

* refactor(rollout): sync-only engine contract; concurrent-safe sglang backends

Drop agenerate/run_session/CoroutineFactory from the rollout-engine contract
(BaseRolloutEngine, BaseSingleTurnRolloutEngine) and make the SGLang backends
safely callable from concurrent threads, so the agentic drain can drive sync
generate from one thread per trajectory:

- native: replace the SessionRunner run_until_complete-under-drive-lock model
  with a serve/park LoopThread over SGLang's own engine.loop — callers submit
  coroutines threadsafe and stay in flight together; weight/memory verbs
  require quiesced generation and run with the loop parked (the Engine's sync
  wrappers still drive the idle loop themselves).
- http: pure sync (urllib + threading.Semaphore + per-batch thread fan-out);
  the httpx client and the backend-owned event loop are gone. Controls are
  bounded 10s best-effort POSTs.
- Backend protocol: generate is sync + thread-safe; async surface removed.

New CPU tests cover the LoopThread lifecycle (concurrent overlap, semaphore
bound, park/serve, quiesce guard, close-waits) and the sync HTTP backend
against a stdlib ThreadingHTTPServer stub.

* refactor(rollout): drop the async surface from the sync-core engines

trainside, vllm_omni, sglang_diffusion, composed: delete agenerate (a
to_thread wrapper over the locked sync path) and run_session + the
LocalAsyncRuntime each held only to serve it; their sync generate paths are
unchanged (the generate lock stays their concurrency story). Drop agenerate
from the RolloutEnginePort protocol (its only consumer, AgentLoop, is sync).
Delete unirl/rollout/engine/runtime.py and its test file; SessionRunner keeps
a local transitional CoroutineFactory alias until the test fakes move off it.

* refactor(agentic): thread-pool drain; sync env step/close; sync test suite

AgenticRolloutEngine now drives trajectories on a per-drive thread pool
(per_worker_concurrency threads = the trajectory cap; a thread holds its
trajectory across tool-wait, preserving the two-bounds design) instead of
coroutines on the inner engine's loop: _drain_worker pulls (blocking ray.get)
and runs a fully-sync _run_one (inner.generate per turn + env.step + duck-typed
env.close teardown). run_drain joins every drain thread and re-raises the first
failure; a failed worker sets _stopping so siblings checkpoint. The barrier
agenerate (test-only) is deleted; coordinator verbs and the turn-boundary
checkpoint/carry semantics are unchanged, and the weight-sync quiesce is now
'abort joined the drain threads => decode-idle'.

Environment protocol: astep/aclose (bridges that existed only for the deleted
loop model) are removed; sync close(sample) joins reset/step as the guaranteed
teardown hook (ToolEnvironment._end_sessions / AlfworldEnv._teardown_episode
bodies, now public). SessionRunner is deleted (backends/base.py is protocol-
only again). Tests: _fakes.py is a sync thread-safe backend/engine pair with a
hold-gate for deterministic overlap; the engine tests assert the same
contracts under threads (cap saturation, two distinct bounds, checkpoint/
resume, conservation, continuous batching); loop tests use step/close, with
thread-based re-entrancy replacing the astep suite.

* docs(rollout): sync-only contract in README + stale async docstring refs

Rewrite the README's generation-interface and extending sections for the
sync-only contract (threads, not asyncio; concurrent-caller requirement for
agentic inners; the loop survives only inside the native SGLang backend).
Update tool/env docstrings that still referenced the deleted astep/aclose/
SessionRunner, and guard the HTTP batch path against an empty wire (the old
gather path returned [] there; a zero-worker pool would raise).

* Fix trainside smoke replay conditions

* Handle ragged AR evaluation batches

* docs: document agentic integration and migration

Add public Sample/Part and agent-loop guides, refresh trainer and recipe documentation, and remove stale retired-API references. Align formatting, executable modes, and pytest discovery so the upstream validation gate is reproducible; cover the merged SGLang top-k behavior.

* chore: remove integration smoke artifacts

* Fix Transformers checkpoint key remapping

* Test model-aware checkpoint remapping

* fix(sampling): require explicit generation frontiers

* fix(agentic): finalize drained rollouts atomically

* fix(agentic): apply safe partial-rollout tail policies

* test: remove tests directory

* fix(sample): resolve interior-stage conditioning at its own frontier

Sample.conditioning is frontier-aligned by contract: turns() gathers one row
per parts[-1] sample. UnifiedModelTrainer._build_request_sample forks
hierarchically into [input, ar(P*N), image(P*N*M)], so HunyuanImage3 t2ti drove
its AR pass at P*N*M width and then filled a P*N-row Part — a deterministic
mismatch whenever the diffusion branch M > 1.

Add Sample.conditioning_at(index), naming the parts[:index+1] replay idiom that
the conditioning()/turns() docstrings already document, and rebuild t2ti on the
shape BagelPipeline t2ti already uses: resolve AR prompts at the AR frontier,
assert the width, fill the AR Part, then re-read conditioning() so the walk
broadcasts prompt and CoT onto the image frontier. Collapse Bagel's hand-rolled
equivalent onto the helper so the two t2ti implementations cannot drift.

* fix(trainside): reject forward_batch_size for multi-stage Samples

The chunk path assumes only parts[-1] is generated: it captures parts[:-1] once
and keeps only chunk.parts[-1]. Given [input, ar, diffusion] a composed pipeline
fills BOTH gen Parts per chunk, so the interior stage is regenerated and then
discarded and the Sample returns carrying its original empty shell. For PE it is
worse — PEPipeline.generate runs both levels inside one call and derives
M = len(diff) // len(ar) from row counts, so slicing diff but not ar either
raises a confusing arity error or silently mis-pairs prompts to images.

Reject the combination instead. The guard keys on gen-Part count, not Part count
(V2V is legitimately 3 Parts / 1 gen Part and sets forward_batch_size: 1), and
fires whenever the knob is set rather than only when this batch would chunk, so
a config cannot pass at one fan-out and corrupt at a larger one. The message
points at per-stage chunking, which pe_sglang_* already does on the SGLang
diffusion sub-engine.

Swept all 113 recipes: 44 pair trainside with forward_batch_size, none of them
multi-gen, so no existing recipe changes behavior.

* fix(agentic): namespace fresh rollout roots per drive

rollout_id alone cannot make root ids unique. _next_batch refills a short buffer
by re-submitting under the SAME rollout_id — up to _MAX_REFILLS get_samples()
calls per rollout step — and a data source may restart its ids on every call
(DefaultDataSource numbers by batch position, not prompt identity). Two drives
then namespace different source rows identically, so _GroupAssembler, which
buckets purely by root id, merges siblings of unrelated prompts into one GRPO
group and _build_tasks overwrites the buffered group's _gt_by_root answer.

Add a monotonic _drive_seq applied to fresh roots only. Carried partials keep
the ids they were submitted under and still rejoin their own siblings.

* fix(agentic): mark infrastructure failures and exclude them from GRPO

_run_one converted every exception into done=True with no failed status, so a
backend outage, tool timeout, or context overflow entered GRPO as a legitimate
low-scoring sibling and biased its whole group. _group_advantages already
documents "NaN reward = crashed trajectory: excluded from the group's mean/std",
but nothing ever produced that NaN, leaving the branch unreachable.

Attach NaN on the failure path (dropping any partial env_reward, which does not
describe a complete trajectory) while still returning done=True so the drain
cannot stall. AgenticEnvTrainer now reports NaN for gen-less trajectories, and
the answer-graded path overwrites graded-but-failed trajectories via a new
_is_failed helper and logs the count.

This intentionally reverses the previous "failed trajectory stays a legit group
member" behaviour: scoring an infrastructure fault as a genuine miss manufactures
a gradient for every sibling in the group.

* fix(agentic): namespace partial rollout roots per drive

Prevent refill batches from reusing root IDs and combining unrelated trajectories in partial-rollout GRPO groups.

* Wire CUDA 13 compatibility into Taiji DRPO launcher

* Source Taiji network environment before strict shell mode

* Make Taiji launcher safe under inherited nounset

* Deduplicate replicated DAPO source rows

* Pin CUDA 13 compiler for SGLang JIT

* Enable deterministic FSDP rollout parity for DRPO

* Expose CUDA runtime to SGLang JIT linker

* fix(veomni): select installed flash attention backend

* fix(distributed): initialize transfer NCCL eagerly

* fix(pe): use expandable CUDA allocator for long runs

* fix(qwen3): bound cached decode growth

* chore(examples): remove task-local launcher defaults

* refactor(agentic): name entrypoints by capability, not benchmark

The six agentic entrypoints were the only ones in the repo named after a
dataset. Every other entrypoint is named for its trainer -- trainer/ar.py ->
train_ar.py, whose docstring states outright that it serves both qwen_vl and
qwen3 recipes. Apply that existing convention:

  train_deep_research.py         -> train_agentic.py
  train_partial_deep_research.py -> train_agentic_partial.py
  train_async_deep_research.py   -> train_agentic_async.py
  train_alfworld.py              -> train_agentic_env.py
  train_partial_alfworld.py      -> train_agentic_env_partial.py
  train_async_alfworld.py        -> train_agentic_env_async.py

Each name now maps 1:1 onto its trainer module, and both axes are legible:
_env is the reward source (environment return vs graded terminal answer),
_partial/_async is the execution topology. Topology is a suffix so the whole
family groups under train_agentic*; train_async_ar.py keeps its prefix rather
than widen the diff.

Recipes stay benchmark-named -- the entrypoint names the capability, the
recipe names the task. Each docstring gains a train_ar.py-style line naming
the recipe family it serves, and the README tables become reward-source x
topology matrices instead of benchmark lists.

Bodies are unchanged: same trainer classes, same constructor arguments, same
config_name defaults. No behavioral change.

Note: python -m unirl.train_alfworld and train_deep_research no longer exist;
six user-facing module paths moved.

* chore: apply ruff-format to qwen3-omni sample-native files

Three files added by adf70e8 were never run through ruff-format, so
`pre-commit run --all-files` -- which the Lint workflow runs on every PR to
main -- fails on this branch. Formatting only; no logic touched.

* fix(agentic): score env-reward faults NaN on every path, not just the barrier

7ad5b34 established that an infrastructure fault must be excluded from GRPO
rather than scored as a genuine miss -- "scoring an infrastructure fault as a
genuine miss manufactures a gradient for every sibling in the group" -- but it
only reached AgenticEnvTrainer. The colocate-partial and fully-async env
variants each carried their own copy of _rewards_and_groups and kept the
pre-fix 0.0, agentic_env_async.py even retaining the comment "gen-less /
failed trajectory stays a legit group member" that the fix says it reverses.

Concretely: when a trajectory faults before its first turn the engine attaches
no reward (_attach_env_reward returns early on an empty gen list), so the
trainer's missing-reward branch decides the sentinel. NaN is dropped from the
group's mean/std and given zero advantage; 0.0 instead drags the mean, shifts
the std, and hands every sibling a gradient the model never earned. A fault
AFTER turn one was already fine on all paths -- the engine attaches NaN there
and it passes straight through.

Rather than patch two literals, hoist the method into a shared
_EnvRewardSource mixin that all three env trainers mix in ahead of their
trainer base. One implementation, so the next sentinel change cannot land on
some paths and miss others. Net -83/+30 in the trainers; the reward SOURCE
becomes a named concept matching how the docstrings already describe it.

Tests, restoring coverage this branch lost:

- tests/rollout/test_agentic_failure_marking.py, deleted by c566479
  "test: remove tests directory" while the fix it guarded stayed. Restored
  verbatim; it still passes against current code.
- tests/trainer/test_agentic_env_reward_source.py, new. The gap that let this
  bug through: nothing ever asserted the env TRAINERS' sentinel. Parametrized
  over all three env paths, plus a structural test that they resolve to one
  shared function.

Verified against the pre-fix tree: exactly 3 failures, naming the two buggy
paths and the divergence itself. 18 pass after the fix.

* test: remove tests directory

Removes the 13 remaining test files (1117 lines), including the two added in
57edcca. No workflow in .github/ runs pytest and the root pyproject declares
no testpaths, so nothing in CI changes.

* refactor(examples): drop the M1 calculator scaffold, port M2 to all topologies

deep_research_calc_mathverify{,_partial,_async}.yaml were bring-up scaffolding:
M1 existed to prove the agentic loop "WITHOUT external services (the existing
CalculatorTool + a rule-based MathVerify reward) ... so M2 only swaps in the
search/visit tools and the LLM-judge reward (config-only)". M2 landed and the
scaffold has no remaining purpose.

The catch: all three were the hydra config_name DEFAULTS for the three
answer-graded entrypoints, and search_judge existed in barrier flavour only.
Deleting them would have left train_agentic_partial and train_agentic_async
with no recipe. Repointing both at the barrier recipe is not an option either:
it declares TensorWeightSync, so the async entrypoint would compose and then
fail at runtime, since disaggregated training needs NCCLWeightSync to cross the
slab boundary.

So port M2 to the two missing topologies, grafting each one's topology block off
the calc recipe it replaces and leaving tools/reward/system-prompt identical:

  deep_research_search_judge_partial.yaml  colocate, TensorWeightSync,
    oversample_batch_size 6, partial_rollout, worker_max_concurrency 24
  deep_research_search_judge_async.yaml    disaggregated, NCCLWeightSync,
    train_fraction 0.5, oversample_batch_size 8, mem_fraction_static 0.8

tail_policy: carry is correct for both -- SearchTool and VisitTool subclass the
stateless Tool, not StatefulTool, so a carried Sample holds all resume state
(the same reason the calc recipes carried).

Verified: all six agentic entrypoints compose against their defaults with
--cfg job --resolve, and the resolved configs carry the right weight-sync class,
mem_fraction_static, and prompts_per_rollout == oversample_batch_size per
topology. check-recipe-targets resolves 2135 paths; hooks and bash -n clean.

Neither new recipe has been train-verified -- they are compose-checked ports.
CalculatorTool itself is left in place; no recipe uses it now, but it stays a
supported library tool.

---------

Co-authored-by: Jianghai <72591262+CjhHa1@users.noreply.github.com>
Co-authored-by: CjhHa1 <cjh18671720497@outlook.com>
Co-authored-by: leviking98z-rgb <leviking98z@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant